import cmd
import sys
import os
import functools

# 使用base code中的common 和 pocketflow
parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(parent_dir)
from base_code import *
from base_code.common import *

# ==============================================================================
# simple_mas_shell 代码介绍
# 在这个案例中，我们引入两个Agent分别处理命令执行和界面显示（交互逻辑仍由cmd库处理）
#  - BDIAgent： 构建命令的工作流
#  - HelpAgent： 显示帮助窗口
#  - CommandRegistry： 存储当前支持的命令信息
# 我们会继续拓展这个案例，来看如何简化这个交互，使用户能充分利用所有已实现工具
# ==============================================================================

# ==============================================================================
# 1. CommandRegistry 是ToolSystem的继承类，添加对支持的命令列表的显示
# ==============================================================================
class CommandRegistry(ToolSystem):
    def __init__(self):
        super().__init__()
        self.lines = [
            "     欢迎使用 SimpleAgent v0.2 (工作流版)",
            " ",
            "在CMD框架下，所有逻辑由Agent的P-D-A循环由工作流处理",
            " ",
            "set <key> <value>        - 设置一个键值对",
            "get <key>                - 获取一个键对应的值",
            "delete <key>             - 删除一个键值对",
            "list                     - 列出所有存储的键值对",
            "use_tool <tool> <param>  - 使用工具",
            "ls path                  - 列出path目录中的文件列表",
            "help                     - 显示帮助信息",
            "exit                     - 退出程序"
        ]


##################################################################
# 全局单例对象：工具系统（环境层）
##################################################################
command_system = CommandRegistry()


##################################################################
# 装饰器：简化工具注册
##################################################################
def tool(name: str, description: str, required_args: List[str], optional_args: Optional[Dict[str, Any]] = None):
    """
    用于注册外部工具函数的装饰器。
    - name: 工具名
    - description: 工具描述
    - required_args: 必须参数列表
    - optional_args: 可选参数字典
    """
    def decorator(func: Callable):
        # 将原函数包装为 Tool 并注册
        new_tool = Tool(
            name=name,
            description=description,
            function=func,
            required_args=required_args,
            optional_args=optional_args
        )
        command_system.register_tool(new_tool)

        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        return wrapper
    return decorator


# ==============================================================================
# 2. HelpAgent 支持UI绘制
# ==============================================================================
class HelpAgent:
    def __init__(self, registry: CommandRegistry):
        self.command_registry = registry

    def show_intro(self):
        try:
            return create_aligned_intro(self.command_registry.lines)
        except Exception:
            return "\n".join(self.command_registry.lines)

    def show_help(self, command=None):
        if command and self.command_registry.has_tool(command):
            t = self.command_registry.get_tool(command)
            usage = f"{t.name} {' '.join(t.required_args)}"
            return f"命令: {t.name}\n功能: {t.description}\n用法: {usage}"
        else:
            return self.show_intro()


# ==============================================================================
# 3. BDIAgent 的核心组件
# ==============================================================================
class AgentMemory:
    """
    Memory class that stores short and long memory of agents

    """
    def __init__(self):
        self.beliefs = {}
        self.command = None
        self.raw_arg = None
        self.desire = None       # (tool_name, kwargs_dict)
        self.error_message = None
        self.result = None
        self.registry = None

    def reset_workflow_data(self):
        """在每次运行前重置临时数据。"""
        self.command, self.raw_arg, self.desire, self.error_message, self.result = (None,) * 5


class PerceptionNode(Node):
    """
    感知节点：解析输入，形成结构化愿望或发现错误。
    """
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        command = (shared.command or "").strip()
        raw_arg = (shared.raw_arg or "").strip()
        registry = shared.registry

        if not registry:
            shared.error_message = "❌ Agent 未连接到 registry"
            return "error"

        if not registry.has_tool(command):
            shared.error_message = f"❌ 未知命令 '{command}'"
            return "error"

        # 规范化 desire：把 memory 和 raw_arg 注入 kwargs，工具内部自行解析 raw_arg
        kwargs = {"memory": shared, "raw_arg": raw_arg}
        shared.desire = (command, kwargs)
        return "success"

    def post(self,shared,prep_res,exec_res):
        return exec_res


class DeliberationNode(Node):
    """
    决策节点：根据愿望选择一个计划（意图），返回对应的动作。
    """
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        # 决策简单：desire 已经包含目标工具与参数，直接走成功路径
        return "success"

    def post(self,shared,prep_res, exec_res):
        return exec_res


class ToolExecutionNode(AsyncNode):
    def __init__(self, tool_system: ToolSystem):
        super().__init__()
        self.tool_system = tool_system

    async def prep_async(self, shared: AgentMemory) -> AgentMemory:
        return shared

    async def exec_async(self, shared: AgentMemory):
        if not shared.desire:
            shared.result = "❌ 无欲望(desire)来执行"
            return

        tool_name, kwargs = shared.desire
        try:
            tool_func = self.tool_system.get_tool(tool_name)
        except KeyError:
            shared.result = f"❌ 工具 '{tool_name}' 未找到"
            return

        # 确保 memory 注入
        if "memory" not in kwargs:
            kwargs["memory"] = shared

        try:
            maybe_coro = tool_func(**kwargs)
            if asyncio.iscoroutine(maybe_coro):
                ret = await maybe_coro
            else:
                ret = maybe_coro
            # 工具函数可以直接把结果写入 memory.result，也可以返回字符串
            if isinstance(ret, str) and (not shared.result):
                shared.result = ret
        except Exception as e:
            shared.result = f"❌ 工具 '{tool_name}' 执行失败: {e}"

        return shared


class ErrorNode(Node):
    """处理流程中产生的错误。"""
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        shared.result = shared.error_message


class FinalizeNode(Node):
    """
    所有成功路径的终点，确保结果被正确设置。
    这个节点目前什么也不做，只是作为一个清晰的汇聚点。
    """
    def prep(self, shared: AgentMemory) -> AgentMemory:
        return shared

    def exec(self, shared: AgentMemory):
        pass


# ==============================================================================
#  基于流程的BDIAgent定义
# ==============================================================================
class BDIAgent(AsyncFlow):
    """
    一个本身就是异步工作流的BDI Agent。
    它继承自AsyncFlow，其流程定义即其内部逻辑。
    """
    def __init__(self):
        # 首先调用父类的构造函数
        super().__init__()

        # 集成 memory 部件
        self.memory = AgentMemory()
        # 引用全局command_system（ToolSystem)
        self.memory.registry = command_system

        # 1. 定义nodes
        perception = PerceptionNode()
        deliberation = DeliberationNode()
        error_handler = ErrorNode()
        finalizer = FinalizeNode()
        executor = ToolExecutionNode(self.memory.registry)

        # 2. 编排工作流 (使用 self.start() 和 self.next())
        self.start(perception)

        perception - "success" >> deliberation
        perception - "error" >> error_handler
        deliberation - "success" >> executor
        executor >> finalizer

        print("💡 BDIAgent (Inheritance-based) 已初始化。Agent本身即是工作流。")

    async def process_command_async(self, command: str, raw_arg: str) -> str:
        """
        Agent的异步入口方法。

        :param command: 命令名
        :param raw_arg: 原始参数字符串
        :return: 执行结果
        """
        # 1. 准备记忆，理论上应放到belief中(sense)，但此处belief主要用于存储用户输入键值
        self.memory.reset_workflow_data()
        self.memory.command = command
        self.memory.raw_arg = raw_arg

        # 2. 执行自身所代表的异步工作流
        #    perception -> ...
        await self.run_async(self.memory)

        # 3. 返回结果
        return self.memory.result



# ==============================================================================
# 注册工具
# ==============================================================================
@tool("set",  "设置键值", ["memory", "raw_arg"])
async def set_value(memory: AgentMemory, raw_arg: str):
    parts = raw_arg.split(None, 1)
    if len(parts) < 2:
        return "❌ 错误: 'set' 命令需要一个键和一个值。\n   用法: set <key> <value>"
    key, value = parts
    memory.beliefs[key] = value
    return f"✅ '{key}' 已设置为 '{value}'"

@tool("get",  "获取键值", ["memory", "raw_arg"])
async def get_value(memory: AgentMemory, raw_arg: str):
    key = raw_arg.strip()
    if not key:
        return "❌ 错误: 'get' 需要一个键。 用法: get <key>"
    if key in memory.beliefs:
        return f"➡️ {key} = {memory.beliefs[key]}"
    else:
        return f"🔍 未找到键: '{key}'"


@tool("list",  "列出键值", ["memory", "raw_arg"])
async def list_beliefs(memory: AgentMemory, raw_arg: str):
    if not memory.beliefs:
        return "ℹ️ 数据库为空"
    else:
        return "\n".join([f"{k}: {v}" for k, v in memory.beliefs.items()])


@tool("delete",  "删除键值", ["memory", "raw_arg"])
async def delete_key(memory: AgentMemory, raw_arg: str):
    key = raw_arg.strip()
    if not key:
        return "❌ 错误: 'delete' 需要一个键。 用法: delete <key>"
    if key in memory.beliefs:
        del memory.beliefs[key]
        return f"✅ 成功: 键 '{key}' 已被删除。"
    else:
        return f"🔍 未找到键: '{key}'"


@tool("ls",  "列出目录", ["memory","raw_arg"])
async def ls_dir(memory: AgentMemory, raw_arg: str = '.'):
    path = raw_arg.strip() or "."
    try:
        if not os.path.isdir(path): return f"错误: '{path}' 不是一个有效的目录。"
        files = os.listdir(path)
        return f"目录 '{path}' 下的内容:\n- " + "\n- ".join(files) if files else f"目录 '{path}' 为空。"
    except Exception as e:
        return f"执行时发生错误: {e}"


@tool("use_tool",  "使用工具", ["memory","raw_arg"])
async def use_tool(memory: AgentMemory, raw_arg: str ):
    arg = raw_arg.strip()
    if not arg:
        return "❌ 用法: use_tool <tool_name> [args]"

    parts = arg.split(None, 1)
    tool_name = parts[0]
    inner_raw = parts[1] if len(parts) > 1 else ""

    if not memory.registry.has_tool(tool_name):
        memory.result = f"❌ 没有找到工具 '{tool_name}'"
        return "error"

    tool_func = command_system.get_tool(tool_name)
    # 调用目标工具，注入 memory 和 inner_raw
    try:
        maybe = tool_func(memory=memory, raw_arg=inner_raw)
        if asyncio.iscoroutine(maybe):
            return await maybe
    except Exception as e:
        memory.result = f"❌ 调用工具失败: {e}"
        return "error"


# =======================================================================
# SimpleAgentShell（应用层）
# =======================================================================
class SimpleAgentShell(cmd.Cmd):
    """
    交互式Shell，通过 asyncio.run() 调用Agent的异步方法完成用户指令。
    """
    help_agent = HelpAgent(command_system)
    cmd_agent = BDIAgent()
    intro = help_agent.show_intro()
    prompt = "agent> "

    def __init__(self):
        super().__init__()

    # --- do_* 方法现在调用 agent 的异步入口 ---
    # 我们使用 asyncio.run() 来从同步的 cmdloop 中调用异步代码。
    # 这会为每个命令创建一个新的事件循环，对于CLI应用来说，简单有效。

    def default(self, line):
        if not line.strip():
            return
        parts = line.split(maxsplit=1)
        cmd_name = parts[0]
        arg = parts[1] if len(parts) > 1 else ""
        result = asyncio.run(self.cmd_agent.process_command_async(cmd_name, arg))
        print(result or (self.cmd_agent.memory.result if self.cmd_agent.memory else ""))

    def do_list_tools(self, arg):
        print("可用命令:", self.cmd_agent.memory.registry.list_tools())

    # --- UI控制相关方法
    def do_exit(self, arg):
        print("\n👋 感谢使用，再见！")
        return True
    def do_quit(self, arg): return self.do_exit(arg)
    def do_EOF(self, arg): print(); return self.do_exit(arg)
    def do_help(self, arg):
        print(self.help_agent.show_help(arg.strip()))

def main():
    """程序主入口。"""
    try:
        SimpleAgentShell().cmdloop()
    except KeyboardInterrupt:
        print("\n👋 检测到中断，程序退出。")
        sys.exit(0)

if __name__ == '__main__':
    main()